﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.Drawing.Imaging;
using System.Linq;
using System.Text;

namespace MyWatermark
{

    //另一种方式
    class DwtTextWatermarking
    {
        // 执行1D离散小波变换（处理偶数长度）
        private static double[] PerformDwt1D(double[] data)
        {
            int n = data.Length;
            if (n % 2 != 0) throw new ArgumentException("Data length must be even.");
            double[] result = new double[n];
            for (int i = 0; i < n / 2; i++)
            {
                result[i] = (data[2 * i] + data[2 * i + 1]) / Math.Sqrt(2);
                result[i + n / 2] = (data[2 * i] - data[2 * i + 1]) / Math.Sqrt(2);
            }
            return result;
        }

        // 执行1D逆离散小波变换
        private static double[] PerformInverseDwt1D(double[] data)
        {
            int n = data.Length;
            double[] result = new double[n];
            for (int i = 0; i < n / 2; i++)
            {
                result[2 * i] = (data[i] + data[i + n / 2]) / Math.Sqrt(2);
                result[2 * i + 1] = (data[i] - data[i + n / 2]) / Math.Sqrt(2);
            }
            return result;
        }

        // 正确实现多级2D DWT分解
        private static double[,] PerformDwt2D(double[,] image, int levels = 2)
        {
            int rows = image.GetLength(0);
            int cols = image.GetLength(1);

            // 检查图像尺寸是否为偶数，如果不是则抛出异常或填充
            if (rows % 2 != 0 || cols % 2 != 0)
            {
                throw new ArgumentException("Image dimensions must be even for DWT.");
            }

            double[,] result = (double[,])image.Clone();

            // 逐级分解
            for (int level = 0; level < levels; level++)
            {
                int currentRows = rows >> level; // 当前层级的行数
                int currentCols = cols >> level; // 当前层级的列数

                // 检查当前层级尺寸是否为偶数
                if (currentRows % 2 != 0 || currentCols % 2 != 0)
                {
                    break; // 如果尺寸为奇数，停止分解
                }

                // 行变换（仅处理当前层级的 LL 区域）
                for (int i = 0; i < currentRows; i++)
                {
                    double[] row = new double[currentCols];
                    for (int j = 0; j < currentCols; j++)
                    {
                        row[j] = result[i, j];
                    }
                    double[] dwtRow = PerformDwt1D(row);
                    for (int j = 0; j < currentCols; j++)
                    {
                        result[i, j] = dwtRow[j];
                    }
                }

                // 列变换（仅处理当前层级的 LL 区域）
                for (int j = 0; j < currentCols; j++)
                {
                    double[] col = new double[currentRows];
                    for (int i = 0; i < currentRows; i++)
                    {
                        col[i] = result[i, j];
                    }
                    double[] dwtCol = PerformDwt1D(col);
                    for (int i = 0; i < currentRows; i++)
                    {
                        result[i, j] = dwtCol[i];
                    }
                }
            }

            return result;
        }
        // 修正逆变换逻辑
        private static double[,] PerformInverseDwt2D(double[,] coeffs, int levels = 2)
        {
            int rows = coeffs.GetLength(0);
            int cols = coeffs.GetLength(1);
            double[,] result = (double[,])coeffs.Clone();

            for (int level = levels - 1; level >= 0; level--)
            {
                int currentSize = rows >> level;
                if (currentSize % 2 != 0) continue;

                // 列逆变换
                for (int j = 0; j < currentSize; j++)
                {
                    double[] col = new double[currentSize];
                    for (int i = 0; i < currentSize; i++)
                        col[i] = result[i, j];
                    col = PerformInverseDwt1D(col);
                    for (int i = 0; i < currentSize; i++)
                        result[i, j] = col[i];
                }

                // 行逆变换
                for (int i = 0; i < currentSize; i++)
                {
                    double[] row = new double[currentSize];
                    Array.Copy(result, i * cols, row, 0, currentSize);
                    row = PerformInverseDwt1D(row);
                    Array.Copy(row, 0, result, i * cols, currentSize);
                }
            }
            return result;
        }

        // 修改嵌入逻辑：处理蓝色通道，定位HL子带
        public static Bitmap EmbedTextWatermark(Bitmap hostImage, string watermarkText, double baseAlpha = 0.1)
        {
            // 转换蓝色通道为二维数组
            double[,] blueChannel = new double[hostImage.Height, hostImage.Width];
            for (int i = 0; i < hostImage.Height; i++)
                for (int j = 0; j < hostImage.Width; j++)
                    blueChannel[i, j] = hostImage.GetPixel(j, i).B / 255.0;

            List<bool> binaryWatermark = TextToBinaryWithHeader(watermarkText);
            int watermarkLength = binaryWatermark.Count;

            // 执行DWT
            double[,] dwt = PerformDwt2D(blueChannel, levels: 2);
            int hlSubbandStart = hostImage.Height / 2;

            // 在HL子带嵌入水印
            int watermarkIndex = 0;
            for (int i = 0; i < hlSubbandStart; i++)
            {
                for (int j = hlSubbandStart; j < hostImage.Width; j++)
                {
                    if (watermarkIndex >= watermarkLength) break;
                    double var = CalculateLocalVariance(dwt, i, j);
                    double alpha = baseAlpha * (1 + Math.Log(1 + var));
                    dwt[i, j] += binaryWatermark[watermarkIndex] ? alpha : -alpha;
                    watermarkIndex++;
                }
            }

            // 逆变换
            double[,] idwt = PerformInverseDwt2D(dwt, levels: 2);

            // 生成含水印图像
            Bitmap watermarked = new Bitmap(hostImage);
            for (int i = 0; i < hostImage.Height; i++)
            {
                for (int j = 0; j < hostImage.Width; j++)
                {
                    Color original = hostImage.GetPixel(j, i);
                    int newB = (int)(Math.Clamp(idwt[i, j], 0, 1) * 255);
                    watermarked.SetPixel(j, i, Color.FromArgb(original.R, original.G, newB));
                }
            }
            return watermarked;
        }

        // 改进提取方法：自动获取长度，定位HL子带
        public static string ExtractTextWatermark(Bitmap watermarkedImage, Bitmap originalImage)
        {
            // 提取蓝色通道
            double[,] wmBlue = new double[watermarkedImage.Height, watermarkedImage.Width];
            double[,] origBlue = new double[originalImage.Height, originalImage.Width];
            for (int i = 0; i < watermarkedImage.Height; i++)
            {
                for (int j = 0; j < watermarkedImage.Width; j++)
                {
                    wmBlue[i, j] = watermarkedImage.GetPixel(j, i).B / 255.0;
                    origBlue[i, j] = originalImage.GetPixel(j, i).B / 255.0;
                }
            }

            // DWT变换
            double[,] wmDwt = PerformDwt2D(wmBlue, levels: 2);
            double[,] origDwt = PerformDwt2D(origBlue, levels: 2);

            // 读取头部确定长度
            List<bool> headerBits = new List<bool>();
            int hlStart = watermarkedImage.Height / 2;
            for (int i = 0; i < 4; i++)
            {
                for (int bit = 0; bit < 8; bit++)
                {
                    int x = i / hlStart;
                    int y = hlStart + (i % hlStart);
                    headerBits.Add(wmDwt[x, y] > origDwt[x, y]);
                }
            }
            int length = BitConverter.ToInt32(BitsToBytes(headerBits), 0);

            // 提取水印数据
            List<bool> watermarkBits = new List<bool>();
            int totalBits = 32 + length * 8;
            for (int i = 4; i < totalBits / 8 + 1; i++)
            {
                int x = i / hlStart;
                int y = hlStart + (i % hlStart);
                if (x >= hlStart || y >= watermarkedImage.Width) break;
                watermarkBits.Add(wmDwt[x, y] > origDwt[x, y]);
            }
            return BinaryToTextWithHeader(watermarkBits);
        }

        // 辅助函数：比特列表转字节数组
        private static byte[] BitsToBytes(List<bool> bits)
        {
            byte[] bytes = new byte[(bits.Count + 7) / 8];
            for (int i = 0; i < bits.Count; i++)
                if (bits[i])
                    bytes[i / 8] |= (byte)(1 << (7 - i % 8));
            return bytes;
        }
        // 将文本转换为二进制序列
        private static List<bool> TextToBinaryWithHeader(string text)
        {
            byte[] bytes = Encoding.UTF8.GetBytes(text);
            List<bool> binaryList = new List<bool>();

            // 添加4字节（32位）长度头
            byte[] lengthHeader = BitConverter.GetBytes(bytes.Length);
            foreach (byte b in lengthHeader)
            {
                binaryList.AddRange(Convert.ToString(b, 2).PadLeft(8, '0').Select(c => c == '1'));
            }

            // 添加文本数据
            foreach (byte b in bytes)
            {
                binaryList.AddRange(Convert.ToString(b, 2).PadLeft(8, '0').Select(c => c == '1'));
            }
            return binaryList;
        }

        // 提取时读取长度头
        private static string BinaryToTextWithHeader(List<bool> binaryList)
        {
            // 读取前32位（4字节）长度头
            byte[] lengthBytes = new byte[4];
            for (int i = 0; i < 32; i += 8)
            {
                string byteStr = string.Join("", binaryList.GetRange(i, 8).Select(b => b ? "1" : "0"));
                lengthBytes[i / 8] = Convert.ToByte(byteStr, 2);
            }
            int length = BitConverter.ToInt32(lengthBytes, 0);

            // 读取后续数据
            List<byte> dataBytes = new List<byte>();
            for (int i = 32; i < 32 + length * 8; i += 8)
            {
                string byteStr = string.Join("", binaryList.GetRange(i, 8).Select(b => b ? "1" : "0"));
                dataBytes.Add(Convert.ToByte(byteStr, 2));
            }
            return Encoding.UTF8.GetString(dataBytes.ToArray());
        }

        // 计算局部方差
        private static double CalculateLocalVariance(double[,] image, int x, int y, int blockSize = 3)
        {
            double sum = 0, sumSq = 0;
            int count = 0;
            for (int i = x - blockSize / 2; i <= x + blockSize / 2; i++)
            {
                for (int j = y - blockSize / 2; j <= y + blockSize / 2; j++)
                {
                    if (i >= 0 && i < image.GetLength(0) && j >= 0 && j < image.GetLength(1))
                    {
                        double val = image[i, j];
                        sum += val;
                        sumSq += val * val;
                        count++;
                    }
                }
            }
            double mean = sum / count;
            return sumSq / count - mean * mean;
        }

    }


    public class test
    {
        public static void test2()
        {


            string originalImagePath = @"D:\ywd\work\ncs\新水印的尝试\blind_watermark\blind_watermark\examples\pic\original.jpg";
            string watermarkedImagePath = @"D:\ywd\work\ncs\新水印的尝试\blind_watermark\blind_watermark\examples\pic\out1111111111111111.jpg";

            Bitmap originalBitmap = new Bitmap(originalImagePath);
            string watermarkText = "Hello, World!";
            double alpha = 0.01;

            // 嵌入水印
            Bitmap watermarkedBitmap = DwtTextWatermarking.EmbedTextWatermark(originalBitmap, watermarkText);
            watermarkedBitmap.Save(watermarkedImagePath, ImageFormat.Jpeg);

            // 提取水印
            Bitmap watermarkedFromFile = new Bitmap(watermarkedImagePath);
            string extractedText = DwtTextWatermarking.ExtractTextWatermark(watermarkedFromFile, originalBitmap);
            Console.WriteLine("提取的文本水印: " + extractedText);
        }
    }
}
